package main

import (
	"bytes"
	"io"
	"log"
	"net"
)

func init() {
	RegisterDbFirewall("db2", Db2FirewallHandle)
}

func Db2FirewallHandle(server net.Conn, client net.Conn, firewalls []*FirewallBlock, exit chan bool) {
	if _, err := readAndWrite(server, client); checkErrorNotEOF(err) {
		exit <- true
		return
	}
	if _, err := readAndWrite(client, server); checkErrorNotEOF(err) {
		exit <- true
		return
	}

	bs, err := readPacket(server)
	if err != nil {
		exit <- true
		return
	}

	var buf bytes.Buffer
	var firewall *FirewallBlock

	nn := 0
	buf.Reset()
	var db []byte
	var username []byte
	for {
		dr, n := parseDRDA(bs[nn:])
		for _, v := range dr.Param {
			//数据库名
			if bytes.Equal(v.Code, []byte{0x21, 0x10}) {
				db = ebc2asc(v.Value[:bytes.IndexByte(v.Value, 0x40)])
				//idex := bytes.IndexByte(db, 0x00)

			} else if bytes.Equal(v.Code, []byte{0x11, 0xa0}) {
				username = ebc2asc(v.Value)
				for _, f := range firewalls {
					if bytes.Equal([]byte(f.Database), db) && bytes.Equal([]byte(f.Username), username) {
						firewall = f
						break
					}
				}
				//if !bytes.Equal(ebc2asc(v.Value), []byte(firewall.ModifyUsername)) {
				//	exit <- true
				//	return
				//}
				//v.Value = asc2ebc([]byte(firewall.Username))
				//reduce := (len(v.Value) - len(firewall.ModifyUsername))
				//v.Length += reduce
				//dr.Length += reduce
				//dr.Length2 += reduce
			} else if bytes.Equal(v.Code, []byte{0x11, 0xa1}) {
				//if !bytes.Equal(ebc2asc(v.Value), []byte(firewall.ModifyPassword)) {
				//	exit <- true
				//	return
				//}
				//v.Value = asc2ebc([]byte(firewall.Password))
				//reduce := (len(v.Value) - len(firewall.ModifyPassword))
				//v.Length += reduce
				//dr.Length += reduce
				//dr.Length2 += reduce
			}
		}
		nn += n
		buf.Write(dr.ToBytes())
		if nn >= len(bs) {
			break
		}
	}
	log.Println(bs)
	log.Println(buf.Len())
	log.Println(buf.Bytes())
	log.Println(firewall)

	//写入最新的串
	if _, err = client.Write(buf.Bytes()); checkErrorNotEOF(err) {
		exit <- true
		return
	}

	if _, err = readAndWrite(client, server); checkErrorNotEOF(err) {
		exit <- true
		return
	}
	//if bs, err := readPacket(client); err != nil {
	//	return
	//}
	////log.Println(string(ebc2asc(bs)))
	//
	//_, err = server.Write(bs)
	//if err != nil {
	//	return
	//}

	go func() {
		for {
			//bs,_ := this.ReadAndWrite(server, client)
			bs, err = readPacket(server)
			if err != nil && err != io.EOF {
				exit <- true
				break
			}
			cpbs := bs[:]
			idex := bytes.Index(bs, []byte{0x24, 0x14})
			if idex > -1 {
				cpbs = bs[idex+7:]
				nullByteIndex := bytes.IndexByte(cpbs, 0xff)
				if nullByteIndex > -1 {
					//过滤掉该请求
					if bytes.Index(cpbs[:nullByteIndex], []byte("drop")) > -1 {
						continue
					}
					log.Println(string((cpbs[:nullByteIndex])))
				}
			}
			log.Println(bs)
			client.Write(bs)

			//_, err = client.Write(bs)
			//if err != nil {
			//	log.Println(err)
			//}
		}
	}()
	go func() {
		for {
			if _, err = readAndWrite(client, server); checkErrorNotEOF(err) {
				exit <- true
				return
			}
		}
	}()
}

//func stepUntil(bs []byte, ptr *int, target []byte) error {
//	if len(target) == 0 {
//		return errors.New("")
//	}
//	for ; *ptr < len(bs); *ptr++ {
//		if target[0] == bs[*ptr]{
//			flag := true
//			for i := 1; i < len(target); i++{
//				flag = flag && (target[i] == bs[i + *ptr])
//			}
//			if flag {
//				*ptr += len(target)
//				break
//			}
//		}
//	}
//	return errors.New("")
//}

type Drda struct {
	Length       int
	Magic        byte
	Format       byte
	ConnectionId []byte
	Length2      int
	CodePoint    []byte
	Param        []*DrdaParameter
}

type DrdaParameter struct {
	Length int
	Code   []byte
	Value  []byte
}

func (dr *Drda) ToBytes() []byte {
	var buf bytes.Buffer
	buf.Write([]byte{byte(dr.Length / 256), byte(dr.Length % 256)})
	buf.WriteByte(dr.Magic)
	buf.WriteByte(dr.Format)
	buf.Write(dr.ConnectionId)
	buf.Write([]byte{byte(dr.Length2 / 256), byte(dr.Length2 % 256)})
	buf.Write(dr.CodePoint)
	for _, v := range dr.Param {
		buf.Write([]byte{byte(v.Length / 256), byte(v.Length % 256)})
		buf.Write(v.Code)
		buf.Write(v.Value)
	}
	return buf.Bytes()
}

func parseDRDA(buf []byte) (dr *Drda, n int) {
	dr = new(Drda)
	n = int(buf[0])*256 + int(buf[1])
	dr.Length = n
	dr.Magic = buf[2]
	dr.Format = buf[3]
	dr.ConnectionId = buf[4:6]
	dr.Length2 = int(buf[6])*256 + int(buf[7])
	dr.CodePoint = buf[8:10]
	ptr := 10
	for ptr < len(buf) && ptr < n {
		param := &DrdaParameter{}
		param.Length = int(buf[ptr])*256 + int(buf[ptr+1])
		param.Code = buf[ptr+2 : ptr+4]
		param.Value = buf[ptr+4 : ptr+param.Length]
		log.Println(string(ebc2asc(param.Value)))
		dr.Param = append(dr.Param, param)
		ptr += param.Length
	}
	return
}

func checkErrorNotEOF(err error) bool {
	return err != nil && err != io.EOF
}
